# 0. 引入必要的包
import time
import os
import numpy as np
import glob
from util import get, preprocess_image, dump
from tqdm import tqdm
from sklearn.model_selection import train_test_split

# 1. 读取配置文件中的信息
train_dir = get('train')  # 获取训练数据路径
char_styles = get("char_styles")  # 获取字符样式列表，注意：必须是列标
new_size = get("new_size")  # 获取新图像大小元组, 注意：必须包含h和w
Xy_root = get('Xy_root')  # 获取Xy文件的存放路径

# 2. 生成X,y
X, y = [], []

print("# 读取训练数据并进行预处理，")
for i in char_styles:
    # 查找所有以 train_{i} 开头的文件
    file_name = glob.glob(f"{train_dir}/train_{i}*")#获取该样式下的所有训练图像文件名。

    # 对于每个样式的图像
    for element in tqdm(file_name, desc=f"处理 {i} 图像：", unit="bit"):
        # 调用 util.py 文件中的 preprocess_image 函数处理每一张图像，获取标签生成标签列
        pie = preprocess_image(element, new_size)
        X.append(pie)#将特征向量和标签分别添加到X和y中。
        # 从文件名中提取标签
        label = os.path.splitext(os.path.basename(element))[0].split("_")[1]
        # 将标签转换为索引
        label = char_styles.index(label)
        y.append(label)#将特征向量和标签分别添加到X和y中。


# 转换为numpy数组
X, y = np.array(X), np.array(y).astype(np.int64)

# 3. 分割测试集和训练集
print("# 将数据按 80% 和 20% 的比例分割")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 5. 序列化分割后的训练和测试样本
dump((X_train, X_test, y_train, y_test), '(X_train, X_test, y_train, y_test)', f'{Xy_root}/Xy')
